#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


###############################################################################
## Trifocal 2D camera location estimation from coplanar points. The
## input to this program are files produced by extract_edgels.


from __future__ import print_function

import sys
import time
import matplotlib
#matplotlib.use('Agg') 

if __name__ == '__main__':
  if sys.argv[0][-7:] == '-nox.py':
    matplotlib.use('Agg') 

from pylab import *
import scipy.io
from numpy import dot
from scipy.optimize import fmin_cg, fmin_powell, fmin
from camori import EdgelSet, PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors, color_distance
from camori_aux import find_line_err, calculate_transformed_coordinates
import simplejson
import Image

set_printoptions(precision=3)

def pplot(x,*a,**kw):
  if len(x.shape)>1:
    plot(x[:,0], x[:,1], *a, **kw)
  else:
    plot(x[0], x[1], *a, **kw)

if __name__ == '__main__':

  ## Plot stuff immediately
  ion()

  rc('text',usetex=False)

  ## Avoid zero divide warnins...
  np.seterr(divide='ignore')

  #PlotStuff=False
  PlotStuff=True
  if sys.argv[0][-7:] == '-nox.py':
    PlotStuff = True



  #################################################################
  ## Load image and initialize pic object

  ## Sets filename from input argument
  if len(sys.argv) < 2:
    print(sys.argv[0], '<job_file.json> [single frame number]')
    raise Exception('Insufficient number of parameters')

  finput = open(sys.argv[1])
  job_params = simplejson.load(finput)
  finput.close()

  fileroot = job_params['root_directory']

  oldgain = job_params['extraction_gain']
  gain = job_params['matching_gain']

  if len(sys.argv) > 2:
    the_frms = [int(sys.argv[2])]
  else:
    the_frms = range(job_params['first_frame'], 1+job_params['last_frame']) 

  for fk,frm in enumerate(the_frms):
    ## Load edgels
    filename = fileroot+'/edgels/edgels-%04d.npz'%(frm)

    Ia = EdgelSet(filename, oldgain, gain)
    ##########
    ### Just to load the image, what a waste
    filename = fileroot+'/frames/'+job_params['filename_format']%frm
    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel
    # Smooth out
    if ("gaussian_smoothing_factor" in job_params.keys() and 
      job_params["gaussian_smoothing_factor"] > 0):
      for c in range(3):
        imr[:,:,c] = scipy.ndimage.gaussian_filter( imr[:,:,c], double(job_params["gaussian_smoothing_factor"]))
    ## Rescales image
    zzz = zeros( (imr.shape[0]/1,imr.shape[1]/1,3) )
    for chan in range(3):
      zzz[:,:,chan]=scipy.ndimage.interpolation.zoom(imr[:,:,chan], 1.0)
    ## Creates picture object
    if not 'projection_model' in job_params:
      model = 'perspective'
    else:
      model = job_params['projection_model']

    if model == 'perspective':
      ## Intrinsic parameters
      focal_distance = job_params['focal_distance_guess']
      p_point = array(job_params['principal_point_guess'])

      pic = PicturePinhole(zzz, focal_distance, p_point[0], p_point[1])

    elif model == 'equirectangular' or model== 'cylindrical_equidistant':
      pic = PictureEqrec(zzz)

    #################

    ## Pick list of "good" lines
    Ia.extract_lines(1000)

    dirs = quaternion_to_matrix(Ia.orientation)


    qq = loadtxt('quaternions.txt')

    if PlotStuff:
      figure(20, figsize=(12.8,7.2))
      suptitle(u'Linhas extraídas e rotuladas', fontweight='bold', fontsize=20)
      subplot(2,2,1)

      imshow(0.75++pic.frame/1024., interpolation='nearest')
      axis([0,1000,500,0])

      subplot(2,2,1)
      for lab in range(3):
        for ll in Ia.all_lines[lab]:
          if np.max(np.sum(np.abs(ll[1:,:]-ll[:-1,:]),1)) > 500:
            continue
          pplot( ll, '-', color=dir_colors[lab], lw=1)

      for lab in range(3):
        for ll in Ia.s_all_lines[lab]:
          subplot(2,2,lab+2)
          pplot( ll * array([1,gain**-1]), '-', color=dir_colors[lab], lw=1)

      # for lab in range(3):
      #   subplot(2,2,lab+2)
      #   axis('equal')

      for lab in range(3):
        subplot(2,2,lab+2)
        if lab==0:
         axis([pi,-pi,-pi/2,pi/2])
        elif lab==1:
         axis([-pi,pi,pi/2,-pi/2])
        elif lab==2:
         axis([-pi,pi,pi/2,-pi/2])

      savefig(fileroot+'/movie3/fff-%04d.png'%frm, dpi=100)


      ## Plot estimated vanishing point directions
      figure(24, figsize=(12.8,7.2))
      suptitle('Predicted edge directions and estimated orientation', fontweight='bold', fontsize=20)
      axes([0.05,0.1,0.40,0.75])
      imshow(pic.frame/260., interpolation='nearest', extent = (-0.5,1000.5,500.5,-0.5))
      pic.plot_vdirs(gca(), 36, quaternion_to_matrix(Ia.orientation))
      axis('equal')
      axis([400,600,450,50])

      for lab,axx in enumerate([
          [0.52,0.5,0.20,0.35],
          [0.77,0.5,0.20,0.35],
          [0.52,0.1,0.20,0.35],
          [0.77,0.1,0.20,0.35]]):
        axes(axx)
        l1=plot(qq[:frm, lab], 'r-', label='Estimate')[0]
        plot(frm-1,qq[frm-1, lab], 'ro')
        l2=plot(qq[:frm,4+lab], 'b--', label='Reference')[0]

        plot([frm-1,frm-1],[-1,1],'k:')
        plot([0, 250],[0,0],'k:')

        xlim(frm-55, frm+5)
        if lab>0:
            ylim(-0.033,0.033)
        else:
            ylim(0.997,1.0)
      figlegend((l1,l2), ('Estimate', 'Reference'), 'lower center', ncol=2 )


      savefig(fileroot+'/movie3/eee-%04d.png'%frm, dpi=100)

      


